// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    env,
    ffi::{CStr, CString},
    fmt::{Debug, Display},
    os::unix::ffi::OsStrExt,
    path::{Path, PathBuf},
};

use cudax::{driver, runtime};
use dashmap::{mapref::one::Ref, DashMap};
use indexmap::IndexSet;
use once_cell::sync::{Lazy, OnceCell};
use tracing::{debug, error};

type CudaResult<T> = Result<T, runtime::cudaError>;
type ModuleEntry<'a> = Ref<'a, PathBuf, CudaModule>;

static GLOBAL_MODULE_CACHE: Lazy<CudaModuleCache> = Lazy::new(CudaModuleCache::new);

fn find_cubin_file<S: AsRef<str>>(module_name: S) -> Option<PathBuf> {
    const ENV_NAME: &str = "XMIG_CUBIN_DIR";
    const DIR_NAME: &str = "cubin";
    const LIB_PATHS: &[&str] = &[
        "/usr/lib64/xmig",
        "/usr/local/lib/xmig",
        "/usr/lib/xmig",
        "/usr/local/lib64/xmig",
    ];

    static CUBIN_DIRS: OnceCell<IndexSet<PathBuf>> = OnceCell::new();

    let cubin_dirs = CUBIN_DIRS.get_or_init(|| {
        let direct_paths = env::var(ENV_NAME)
            .ok()
            .map(PathBuf::from)
            .into_iter()
            .chain(LIB_PATHS.iter().map(PathBuf::from));
        let relative_paths = env::current_dir()
            .ok()
            .into_iter()
            .chain(
                env::current_exe()
                    .ok()
                    .and_then(|p| p.parent().map(|d| d.to_path_buf())),
            )
            .map(|base| base.join(DIR_NAME));
        IndexSet::from_iter(
            direct_paths
                .chain(relative_paths)
                .filter(|path| path.is_dir())
                .filter_map(|path| path.canonicalize().ok()),
        )
    });

    cubin_dirs.iter().find_map(|dir| {
        let file_path = dir.join(module_name.as_ref());
        file_path.is_file().then_some(file_path)
    })
}

#[inline]
pub fn get_or_load_function<S, T>(module_name: S, func_name: T) -> CudaResult<driver::CUfunction>
where
    S: AsRef<str>,
    T: AsRef<CStr>,
{
    let module_path = match find_cubin_file(module_name) {
        Some(path) => path,
        None => return Err(driver::cudaError_enum_CUDA_ERROR_FILE_NOT_FOUND),
    };
    debug!("module_path: {}", module_path.display());

    GLOBAL_MODULE_CACHE
        .get_or_load_module(module_path)?
        .get_or_load_function(func_name)
}

#[derive(Debug)]
struct CudaModuleCache {
    hmods: DashMap<PathBuf, CudaModule>,
}

impl CudaModuleCache {
    #[inline]
    fn new() -> Self {
        Self {
            hmods: DashMap::new(),
        }
    }

    fn get_or_load_module<S>(&self, module_path: S) -> CudaResult<ModuleEntry<'_>>
    where
        S: AsRef<Path>,
    {
        let module_path = module_path.as_ref();

        if let Some(entry) = self.hmods.get(module_path) {
            return Ok(entry);
        }

        let entry = self
            .hmods
            .entry(module_path.to_owned())
            .or_try_insert_with(|| CudaModule::load(module_path))?
            .downgrade();

        Ok(entry)
    }
}

#[derive(Debug)]
struct CudaModule {
    path: PathBuf,
    hmod: driver::CUmodule,
    hfuncs: DashMap<CString, driver::CUfunction>,
}

unsafe impl Send for CudaModule {}
unsafe impl Sync for CudaModule {}

impl CudaModule {
    fn load<S>(module_path: S) -> CudaResult<Self>
    where
        S: AsRef<Path>,
    {
        let path = module_path.as_ref();
        debug!("Loading cuda module '{}' ...", path.display());

        let fname = CString::new(path.as_os_str().as_bytes())
            .map_err(|_| runtime::cudaError_cudaErrorInvalidValue)?;

        let mut module: driver::CUmodule = std::ptr::null_mut();
        let ret = unsafe { driver::cuModuleLoad(&mut module, fname.as_ptr()) };
        if ret != 0 {
            error!(
                "Failed to load cuda module '{}', ret={}.",
                path.display(),
                ret
            );
            return Err(ret);
        }

        debug!("Loaded cuda module '{}'.", path.display());
        Ok(Self {
            path: path.to_owned(),
            hmod: module,
            hfuncs: DashMap::new(),
        })
    }

    fn get_or_load_function<S>(&self, func_name: S) -> CudaResult<driver::CUfunction>
    where
        S: AsRef<CStr>,
    {
        let func_name = func_name.as_ref();

        if let Some(entry) = self.hfuncs.get(func_name) {
            return Ok(*entry.value());
        }

        let entry = self
            .hfuncs
            .entry(func_name.to_owned())
            .or_try_insert_with(|| {
                debug!(
                    "Loading cuda function '{}' in '{}'...",
                    func_name.to_string_lossy(),
                    self
                );

                let mut hfunc: driver::CUfunction = std::ptr::null_mut();

                let ret = unsafe {
                    driver::cuModuleGetFunction(&mut hfunc, self.hmod, func_name.as_ptr())
                };
                if ret != 0 {
                    error!(
                        "Failed to load cuda function '{}' in '{}', ret={}.",
                        func_name.to_string_lossy(),
                        self,
                        ret
                    );
                    return Err(ret);
                }

                debug!("Loaded cuda function '{}'.", func_name.to_string_lossy());
                Ok(hfunc)
            })?;

        Ok(*entry.value())
    }
}

impl Display for CudaModule {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.path.fmt(f)
    }
}

impl Drop for CudaModule {
    fn drop(&mut self) {
        debug!("Unloading cuda module '{}'...", self);
        let ret = unsafe { driver::cuModuleUnload(self.hmod) };
        if ret != 0 {
            error!("Failed to unload cuda module '{}', ret={}.", self, ret);
        } else {
            debug!("Unloaded cuda module '{}'.", self);
        }
    }
}
